#include <wdm.h>

#pragma warning( disable : 4100 )

VOID KEXCB_Unload(_In_ struct _DRIVER_OBJECT *DriverObject);
NTSTATUS KEXCB_DispatchCreate(_Inout_ struct _DEVICE_OBJECT *DeviceObject,
							  _Inout_ struct _IRP *Irp);
NTSTATUS KEXCB_DispatchClose(_Inout_ struct _DEVICE_OBJECT *DeviceObject,
							 _Inout_ struct _IRP *Irp);
NTSTATUS KEXCB_DispatchRead(_Inout_ struct _DEVICE_OBJECT *DeviceObject,
							_Inout_ struct _IRP *Irp);
NTSTATUS KEXCB_DispatchWrite(_Inout_ struct _DEVICE_OBJECT *DeviceObject,
							 _Inout_ struct _IRP *Irp);

#define DRIVER_NAME "KEXCB"
#define NT_DEVICE_NAME L"\\Device\\KEXCB"
#define DOS_DEVICE_NAME L"\\DosDevices\\KEXCB"

#define EXCHANGE_BUFFER_SIZE 1024
char exchange_buffer[EXCHANGE_BUFFER_SIZE];

/*
http://msdn.microsoft.com/en-us/library/windows/hardware/ff544113%28v=vs.85%29.aspx
*/
DRIVER_INITIALIZE DriverEntry;
NTSTATUS DriverEntry(_In_ struct _DRIVER_OBJECT *DriverObject,
					 _In_ PUNICODE_STRING RegistryPath)
{
	/* Using NTSTATUS values:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff565436%28v=vs.85%29.aspx
	*/
	NTSTATUS status = STATUS_SUCCESS;
	UNICODE_STRING DeviceName;
	UNICODE_STRING DosDeviceName;
	PDEVICE_OBJECT DeviceObject = NULL;

	DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_TRACE_LEVEL,
				DRIVER_NAME": --> DriverEntry\n");
	
	/* Setup driver object */
	DriverObject->MajorFunction[IRP_MJ_CREATE] = KEXCB_DispatchCreate;
	DriverObject->MajorFunction[IRP_MJ_CLOSE] = KEXCB_DispatchClose;
	DriverObject->MajorFunction[IRP_MJ_READ] = KEXCB_DispatchRead;
	DriverObject->MajorFunction[IRP_MJ_WRITE] = KEXCB_DispatchWrite;

	DriverObject->DriverUnload = KEXCB_Unload;

	/* Create the device */
	RtlInitUnicodeString(&DeviceName, NT_DEVICE_NAME);
	RtlInitUnicodeString(&DosDeviceName, DOS_DEVICE_NAME);
	/*
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff548397%28v=vs.85%29.aspx
	*/
	status = IoCreateDevice(DriverObject, 0, &DeviceName,
							FILE_DEVICE_UNKNOWN,
							FILE_DEVICE_SECURE_OPEN,
							FALSE,
							&DeviceObject);
	if (!NT_SUCCESS(status))
	{
		/*
		http://msdn.microsoft.com/en-us/library/windows/hardware/ff543632%28v=vs.85%29.aspx
		*/
		//DbgPrint("Failed to IoCreateDevice\n");
		/*
		http://msdn.microsoft.com/en-us/library/windows/hardware/ff543634%28v=vs.85%29.aspx
		*/
		DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL,
					DRIVER_NAME": Error creating device (Failed to IoCreateDevice)\n");
		return status;
	}

	// Map the "DOS Device Name" to the "NT Device Name"
	status = IoCreateSymbolicLink(&DosDeviceName, &DeviceName);
	if (!NT_SUCCESS(status))
	{
		DbgPrintEx(DPFLTR_IHVDRIVER_ID,	DPFLTR_ERROR_LEVEL,
					DRIVER_NAME": Error creating symlink (Failed to IoCreateSymbolicLink)\n");
		return status;
	}

	// Initializing a Device Object
	// http://msdn.microsoft.com/en-us/library/windows/hardware/ff547807(v=vs.85).aspx
	DeviceObject->Flags |= DO_BUFFERED_IO;
	// must be cleared if DeviceObject was created outside of the DriverEntry function
	// DeviceObject->Flags &= (~DO_DEVICE_INITIALIZING);

	KdPrintEx((DPFLTR_IHVDRIVER_ID, DPFLTR_TRACE_LEVEL,
				DRIVER_NAME": <-- DriverEntry\n"));
	
	return status;
}

/*
http://msdn.microsoft.com/en-us/library/windows/hardware/ff564886%28v=vs.85%29.aspx
*/
DRIVER_UNLOAD KEXCB_Unload;
VOID KEXCB_Unload(_In_ struct _DRIVER_OBJECT *DriverObject)
{
	NTSTATUS status = STATUS_SUCCESS;
	UNICODE_STRING DosDriverName;

	KdPrintEx((DPFLTR_IHVDRIVER_ID, DPFLTR_TRACE_LEVEL,
				DRIVER_NAME": --> Unload\n"));

	RtlInitUnicodeString(&DosDriverName, DOS_DEVICE_NAME);
	status = IoDeleteSymbolicLink(&DosDriverName);

	IoDeleteDevice(DriverObject->DeviceObject);

	DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_TRACE_LEVEL,
				DRIVER_NAME": <-- Unload\n");
}

DRIVER_DISPATCH KEXCB_DispatchCreate;
NTSTATUS KEXCB_DispatchCreate(_Inout_ struct _DEVICE_OBJECT *DeviceObject,
							  _Inout_ struct _IRP *Irp)
{
	/* Controlling Device Namespace Access:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff542068(v=vs.85).aspx
	fail all IRP_MJ_CREATE requests that specify
	an IrpSp->FileObject->FileName parameter whose length is nonzero,
	because we don't support namespaces
	*/
	return STATUS_SUCCESS;
}

DRIVER_DISPATCH KEXCB_DispatchClose;
NTSTATUS KEXCB_DispatchClose(_Inout_ struct _DEVICE_OBJECT *DeviceObject,
							 _Inout_ struct _IRP *Irp)
{
	return STATUS_SUCCESS;
}

/*
http://msdn.microsoft.com/en-us/library/windows/hardware/ff543376%28v=vs.85%29.aspx
*/
DRIVER_DISPATCH KEXCB_DispatchRead;
NTSTATUS KEXCB_DispatchRead(_Inout_ struct _DEVICE_OBJECT *DeviceObject,
							_Inout_ struct _IRP *Irp)
{
	NTSTATUS status = STATUS_SUCCESS;
	PIO_STACK_LOCATION stack = NULL;
	LARGE_INTEGER offset;
	ULONG length;
	PVOID buffer;
	ULONG bytes_to_read;

	stack = IoGetCurrentIrpStackLocation(Irp);

	offset = stack->Parameters.Read.ByteOffset;
	length = stack->Parameters.Read.Length;
	buffer = Irp->AssociatedIrp.SystemBuffer;

	/* Perform the request */
	if (offset.QuadPart < EXCHANGE_BUFFER_SIZE)
	{
		/* read as much bytes as we can */
		bytes_to_read = (offset.QuadPart + length <= EXCHANGE_BUFFER_SIZE)
							? length
							: EXCHANGE_BUFFER_SIZE - (ULONG)offset.QuadPart;
		RtlCopyMemory(buffer, exchange_buffer + offset.QuadPart, bytes_to_read);
	}
	else
	{
		status = STATUS_UNSUCCESSFUL;
		bytes_to_read = 0;
	}

	/* Complete the request*/
	/* IO_STATUS_BLOCK structure:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff550671%28v=vs.85%29.aspx
	*/
	Irp->IoStatus.Status = status;
	Irp->IoStatus.Information = bytes_to_read;

	/* IoCompleteRequest routine:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff548343%28v=vs.85%29.aspx
	*/
	IoCompleteRequest(Irp, IO_NO_INCREMENT);

	return STATUS_SUCCESS;
}

/* Assuming BUFFERED_IO */
DRIVER_DISPATCH KEXCB_DispatchWrite;
NTSTATUS KEXCB_DispatchWrite(_Inout_ struct _DEVICE_OBJECT *DeviceObject,
							 _Inout_ struct _IRP *Irp)
{
	/*
	IRP:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff550694%28v=vs.85%29.aspx
	IO_STACK_LOCATION:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff550659%28v=vs.85%29.aspx
	*/
	NTSTATUS status = STATUS_SUCCESS;
	PIO_STACK_LOCATION stack = NULL;
	LARGE_INTEGER offset;
	ULONG length;
	PVOID buffer;
	ULONG bytes_to_write;
	
	/* Get request parameters */
	/* IoGetCurrentIrpStackLocation routine:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff549174%28v=vs.85%29.aspx
	*/
	stack = IoGetCurrentIrpStackLocation(Irp);

	offset = stack->Parameters.Write.ByteOffset;
	length = stack->Parameters.Write.Length;
	buffer = Irp->AssociatedIrp.SystemBuffer;

	/* Perform the request */
	if (offset.QuadPart < EXCHANGE_BUFFER_SIZE)
	{
		/* write as much bytes as we can */
		bytes_to_write = (offset.QuadPart + length <= EXCHANGE_BUFFER_SIZE)
							? length
							: EXCHANGE_BUFFER_SIZE - (ULONG)offset.QuadPart;
		RtlCopyMemory(exchange_buffer + offset.QuadPart, buffer, bytes_to_write);
	}
	else
	{
		status = STATUS_UNSUCCESSFUL;
		bytes_to_write = 0;
	}

	/* Complete the request*/
	/* IO_STATUS_BLOCK structure:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff550671%28v=vs.85%29.aspx
	*/
	Irp->IoStatus.Status = status;
	Irp->IoStatus.Information = bytes_to_write;

	/* IoCompleteRequest routine:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff548343%28v=vs.85%29.aspx
	*/
	IoCompleteRequest(Irp, IO_NO_INCREMENT);

	return status;
}
